from langchain.tools import tool, ToolRuntime
from pydantic import BaseModel, Field
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

from model.Ark import ArkModel

def base_output():
    model = ArkModel().model
    @tool
    def get_weather(add: str, runtime: ToolRuntime):
        """ 获取一个地址的天气信息
            :Args:
                add: 地理位置
        """
        return '晴天'

    ## 基于模型的 schema
    class CustomerResponseSchema(BaseModel):
        address: str = Field(description='地理位置')
        weather: str = Field(description='天气情况')

    ## 基于json 的schema
    # CustomerJsonSchema = {
    #     "type": "object",
    #     "description": "天气信息的描述",
    #     "properties": {
    #         "address": {"type": "string", "description": "地理位置"},
    #         "weather": {"type": "string", "description": "天气情况"},
    #     },
    #     "required": ["address", "weather"]
    # }

    agent = create_agent(
        model = model,
        tools=[get_weather],
        checkpointer=InMemorySaver(),
        response_format=CustomerResponseSchema
    )

    response = agent.invoke({"messages": [{ "role": 'user', 'content': '今天的天气怎么样？' }]}, { "configurable": { 'thread_id': '12312' } })
    print(response['messages'][-1].content)

def run():
    base_output()
